/*
 * gbdt_tree.cpp
 *
 *  Created on: 2012-7-21
 *      Author: wanghaox
 */

#include <queue>
#include "gbdt_tree.h"

bool GbdtTree::init(int depth) {
	// initialize parameters
	_depth = depth;
	_layer = 0;
	_nodes = 1;

	// initialize levels
	_layers = new TreeNode**[depth];
	for (int i = 0; i < _depth; i++) {
		// initialize layer of nodes
		int n = (int) pow(2.f, (float) i);
		_layers[i] = new TreeNode*[n];

		// initialize nodes
		for (int j = 0; j < n; j++) {
			_layers[i][j] = new TreeNode();
			_layers[i][j]->reset();
			_layers[i][j]->level = i;
			_layers[i][j]->seq = j;
		}
	}

	return true;
}

void GbdtTree::uninit() {
	// delete histograms and nodes
	for (int i = 0; i < _depth; i++) {
		delete[] _layers[i];
		_layers[i] = NULL;
	}

	// delete layers
	delete[] _layers;
	_depth = 0;
	_layer = 0;
	_nodes = 0;
}

void GbdtTree::clear() {
	// clear each node
	for (int d = 0; d < _depth; d++) {
		int n = nodes_at_depth(d);
		for (int j = 0; j < n; j++) {
			_layers[d][j]->reset();
		}
	}
	// reset layer
	_layer = 0;
	_nodes = 1;
}

void GbdtTree::find_local_split(GbdtParam &global_param, GbdtFeature &feature) {
	// recompute counts at nodes
	int *treenode = feature._instance_to_treenode;
	double *residual = global_param._instance_residuals;
	for (int i = 0; i < global_param._instance_num; i++) {
		int node = treenode[i];
		_layers[_layer][node]->m_infty += 1;
		_layers[_layer][node]->l_infty += residual[i];
	}

// set labels at nodes
// (this is redundant for many internal nodes --
// overwriting the existing value with the same value --
// but necessary for the root node and nodes that have stopped short)
	for (int n = 0; n < _nodes; n++) {
		TreeNode* node = _layers[_layer][n];
		TreeNode* child1 = _layers[_layer + 1][n * 2];
		TreeNode* child2 = _layers[_layer + 1][n * 2 + 1];
		if (node->m_infty > 0) {
			double label = node->l_infty / (double) node->m_infty;
			node->label = label;
			child1->label = label;
			child2->label = label;
		}
	}

// iterate over features and update best splits at each node
	update_split(global_param, feature);
}

void GbdtTree::update_split(GbdtParam &global_param, GbdtFeature &feature) {
	// compute global feature index
	int feature_index = feature._global_feature_index;

	// reset counts at nodes
	for (int n = 0; n < _nodes; n++) {
		TreeNode* node = _layers[_layer][n];
		node->m_s = 0;
		node->l_s = 0.0;
	}

	// iterate over feature
	int instance_num = global_param._instance_num;
	double *sored_feature = feature._sored_feature;
	int *sorted_index = feature._sorted_instance;
	int *instrance_nodes = feature._instance_to_treenode;
	double *residual = global_param._instance_residuals;

	for (int j = 0; j < instance_num; j++) {
		// get current value
		double value = sored_feature[j];
		int index = sorted_index[j];
		int instance_node = instrance_nodes[index];
		double l = residual[j];

		// get node
		TreeNode* node = _layers[_layer][instance_node];

		// if not first instance at node and greater than split point, consider new split at v
		if (node->m_s > 0 && value > node->s) {
			double loss_i = pow(node->l_s, 2.0) / (double) node->m_s
					+ pow(node->l_infty - node->l_s, 2.0) / (double) (node->m_infty - node->m_s);
			if (node->loss < 0 || loss_i > node->loss) {
				node->loss = loss_i;
				node->feature = feature_index;
				node->split = (node->s + value) / 2.f;

				// TODO : create a lookup table for these child values at tree construction, store in static tree or store each in static node
				TreeNode* child1 = _layers[_layer + 1][2 * instance_node];
				child1->label = node->l_s / (double) node->m_s;
				TreeNode* child2 = _layers[_layer + 1][2 * instance_node + 1];
				child2->label = (node->l_infty - node->l_s) / (double) (node->m_infty - node->m_s);

				// if (child2->label > 5.0)
				//     printf("### %f %d %d %f %d %f %d %f %d %d %d %f %d\n", v, i, n, l, node->m_s, node->l_s, node->m_infty, node->l_infty, globalf, f, j, loss_i, layer);
			}
		}

		// update variables
		node->m_s += 1;
		node->l_s += l;
		node->s = value;
	}
}

void GbdtTree::print_node(int level, int i, double learning_rate) {
	TreeNode *node = _layers[level][i];

	// print node
	if (level > 0)
		printf(",");
	printf("%d:%f:%f", node->feature, node->split, learning_rate * node->label);

	// print children
	if (node->feature > 0) { // a splitting node
		print_node(level + 1, 2 * i, learning_rate); // print left child
		print_node(level + 1, 2 * i + 1, learning_rate); // print right child
	}
}

void GbdtTree::print() {
	print_node(0, 0, 1.0f);
	printf("\n");
}

void GbdtTree::print_screen() {
	queue<TreeNode *> print_queue;
	TreeNode *root = _layers[0][0];

	print_queue.push(root);
	while (!print_queue.empty()) {
		TreeNode *cur = print_queue.front();
		print_queue.pop();
		if (cur->feature > 0) {
			print_queue.push(_layers[cur->level + 1][cur->seq * 2]);
			print_queue.push(_layers[cur->level + 1][cur->seq * 2 + 1]);
			printf("%d:%f:%f ", cur->feature, cur->split,  cur->label);
		}
	}
	printf("\n");
}

void GbdtTree::curren_layer_to_buffer(string &buffer) {
    // instantiate buffer
    int buffersize = nodes*5;
    double* buffer = new double[buffersize];

    // write layer of tree to buffer
    for (int n=0; n<_nodes; n++) {
        StaticNode* node = _layers[_layer][n];
        buffer[n*5 + 0] = node->loss;
        buffer[n*5 + 1] = node->feature;
        buffer[n*5 + 2] = node->split;

        StaticNode* child1 = layers[layer+1][n*2];
        buffer[n*5 + 3] = child1->label;

        StaticNode* child2 = layers[layer+1][n*2+1];
        buffer[n*5 + 4] = child2->label;
    }

    // get myid and numprocs
    int myid;
    MPI_Comm_rank(MPI_COMM_WORLD, &myid);
    int numprocs;
    MPI_Comm_size(MPI_COMM_WORLD, &numprocs);

    // determine isRoot
    int root = numprocs-1;
    bool isRoot = (myid == root);

    // exchange buffers
    double* rbuffer = (isRoot ? new double[numprocs*buffersize] : NULL);
    MPI_Gather(buffer, buffersize, MPI_DOUBLE, rbuffer, buffersize, MPI_DOUBLE, root, MPI_COMM_WORLD);

    // save best global splits
    if (isRoot) for (int n=0; n<nodes; n++) {
        // reset loss at node and get pointers
        StaticNode* node = layers[layer][n];
        node->loss = -1;
        StaticNode* child1 = layers[layer+1][n*2];
        StaticNode* child2 = layers[layer+1][n*2+1];

        // consider loss from all processors
        for (int p=0; p<numprocs; p++) {
            int offset = p*buffersize + n*5;
            double loss = rbuffer[offset + 0];

            // update if better than current
            if (node->loss < 0 or loss > node->loss) {
                node->loss = loss;
                node->feature = (int) rbuffer[offset + 1];
                node->split = (float) rbuffer[offset + 2];
                child1->label = rbuffer[offset + 3];
                child2->label = rbuffer[offset + 4];
            }
        }
    }

    // buffer best global splits
    if (isRoot) for (int n=0; n<nodes; n++) {
        StaticNode* node = layers[layer][n];
        buffer[n*5 + 0] = node->loss;
        buffer[n*5 + 1] = node->feature;
        buffer[n*5 + 2] = node->split;

        StaticNode* child1 = layers[layer+1][n*2];
        buffer[n*5 + 3] = child1->label;

        StaticNode* child2 = layers[layer+1][n*2+1];
        buffer[n*5 + 4] = child2->label;
    }

    // broadcast best splits
    MPI_Bcast(buffer, nodes*5, MPI_DOUBLE, root, MPI_COMM_WORLD);

    // update tree with best global splits
    for (int n=0; n<nodes; n++) {
        StaticNode* node = layers[layer][n];
        node->loss = buffer[n*5 + 0];
        node->feature = (int) buffer[n*5 + 1];
        node->split = (float) buffer[n*5 + 2];

        StaticNode* child1 = layers[layer+1][n*2];
        child1->label = buffer[n*5 + 3];

        StaticNode* child2 = layers[layer+1][n*2+1];
        child2->label = buffer[n*5 + 4];
    }

    // delete buffers
    delete [] buffer;
    delete [] rbuffer;
}

//bool StaticTree::containsSplittingFeature(FeatureData* data) {
//    int feature; float split;
//    for (int i=0; i<nodes; i++) {
//        getSplit(i, feature, split);
//        if (data->isLocalFeature(feature)) return true;
//    }
//    return false;
//}
//
//void StaticTree::updateTrainingPredictions(FeatureData *data, double learningrate) {
//    int N = data->getN();
//    for (int i=0; i<N; i++) {
//        int node = data->getNode(i);
//        double pred = learningrate * layers[layer][node]->label;
//        data->updatePred(i,pred);
//    }
//}
//
//void StaticTree::updatePredictions(InstanceData *data, double learningrate) {
//    int N = data->getN();
//    for (int i=0; i<N; i++) {
//        double pred = learningrate * classifyDataPoint(data, i);
//        data->updatePred(i,pred);
//    }
//}
//
//double StaticTree::classifyDataPoint(InstanceData* data, int p) {
//	// descend tree
//	int node = 0;
//	for (int i=0; i<depth-1; i++) {
//		// get feature and split point
//		int f = layers[i][node]->feature;
//		float s = layers[i][node]->split;
//
//		// check for valid split, otherwise return
//		if (f < 0) return layers[i][node]->label;
//
//		// perform split
//        node <<= 1; // node *= 2, index of left child
//        node |= (data->getFeature(f,p) >= s); // node += 1, if right child
//	}
//
//	// return label of leaf node as prediction
//	return layers[depth-1][node]->label;
//}
//
//void StaticTree::printNode(int level, int i, double learningrate) {
//    // get node
//    StaticNode *node = layers[level][i];
//
//    // print node
//    if (level > 0) printf(",");
//    printf("%d:%f:%f", node->feature, node->split, learningrate * node->label);
//
//    // print children
//    if (node->feature > 0) { // a splitting node
//        printNode(level+1, 2*i, learningrate); // print left child
//        printNode(level+1, 2*i+1, learningrate); // print right child
//    }
//}
//
//void StaticTree::printTree(double learningrate) {
//    printNode(0,0,learningrate);
//    printf("\n");
//
//    // for (int d=0; d<depth; d++) {
//    //     for (int i=0; i<nodesAtDepth(d); i++) {
//    //         StaticNode *node = layers[d][i];
//    //         printf("%d:%f:%f:%f:%d ", node->feature, node->split, node->label, node->l_infty, node->m_infty);
//    //     }
//    //     printf("\n");
//    // }
//    // printf("\n");
//}
//
//void StaticTree::getSplit(int node, int &feature, float &split) {
//	feature = layers[layer][node]->feature;
//	split = layers[layer][node]->split;
//}
//
//int StaticTree::getNumNodes() {
//    return nodes;
//}

